﻿using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.IO;
using System.IO.Compression;
using System.Security.Cryptography;
using System.Text;
using Util.IO;
using Util.IO.Compression;

namespace GJQT.ZPK
{
#if false
    /// Multiple ZPK volumes can be aggregated to form a ZPK archive with updated files.
    public class ZpkArchiveSet
    {
    }
#endif

    /// <summary>
    /// Provides methods to read a ZPK archive.
    /// </summary>
    public class ZpkArchive
    {
        private string m_filename;
        //private int m_version;      // Version of the volume. This is taken as the
        // version of all the entries contained in
        // this volume.
        private TinyCipher tea;
        private ExtendedTinyCipher xtea;
        //private Utils.TinyCipher m_tea = new Utils.TinyCipher();
        //private Utils.TinyCipherEx m_xtea = new Utils.TinyCipherEx();
        private ReadOnlyCollection<ZpkArchiveEntry> m_entries;

        /// <summary>
        /// Opens a ZPK archive for reading.
        /// </summary>
        /// <param name="fileName"></param>
        public ZpkArchive(string fileName)
        {
            byte[] buffer = new byte[65536];
            List<ZpkArchiveEntry> entries = new List<ZpkArchiveEntry>();

            // Open the file.
            using (FileStream stream = new FileStream(
                fileName, FileMode.Open, FileAccess.Read, FileShare.Read))
            using (BinaryReader reader = new BinaryReader(stream))
            {
                // Read encryption key.
                int keyLen = reader.ReadUInt16();
                reader.ReadFull(buffer, 0, keyLen);

                // Initialize ciphers.
                BlowfishCipher blowfish = new BlowfishCipher(buffer, 0, keyLen);
                tea = new TinyCipher(buffer, 0, keyLen);
                xtea = new ExtendedTinyCipher(buffer, 0, keyLen);

                // Read Index Section offset and seek there.
                int indexSectionOffset = reader.ReadInt32();
                stream.Seek(indexSectionOffset, SeekOrigin.Begin);

                // Read index entries block by block.
                while (true)
                {
                    // Read and decode block length.
                    reader.ReadFull(buffer, 0, 4);
                    int blocklen = ((buffer[0] ^ buffer[3]) << 8) | (buffer[1] ^ buffer[2]);
                    if (blocklen == 0)
                        break;

                    // Read block data and decrypt it.
                    reader.ReadFull(buffer, 0, blocklen);
                    blowfish.TransformBlock(buffer, 0, blocklen, buffer, 0);

                    // Parse each filename in the block.
                    using (MemoryStream blockStream = new MemoryStream(buffer, 0, blocklen))
                    using (BinaryReader blockReader = new BinaryReader(blockStream))
                    {
                        // Since the block cipher works in blocks of 8 bytes,
                        // there may be garbage padding at the end of a 
                        // segment. Such padding must be ignored.
                        while (blockStream.Position <= blockStream.Length - 8)
                        {
                            int nameLen = blockReader.ReadUInt16();
                            //if (nameLen == 0)
                            //    break;

                            ZpkArchiveEntry ent = new ZpkArchiveEntry();
                            ent.FullName = Encoding.GetEncoding("gbk").GetString(
                                buffer, (int)blockStream.Position, nameLen);
                            blockStream.Seek(nameLen, SeekOrigin.Current);

                            ent.EncryptionMethod = (ZpkEncryptionMethod)blockReader.ReadByte();
                            ent.CompressionMethod = (ZpkCompressionMethod)blockReader.ReadByte();
                            ent.CompressedOffset = blockReader.ReadUInt32();
                            ent.Length = blockReader.ReadUInt32();
                            ent.CompressedLength = blockReader.ReadUInt32();
                            ent.Archive = this;
                            entries.Add(ent);
                        }
                    }
                }
            }
            m_filename = Path.GetFullPath(fileName);
            m_entries = new ReadOnlyCollection<ZpkArchiveEntry>(entries);
        }

        /// <summary>
        /// Gets the collection of entries that are currently in the zpk
        /// archive.
        /// </summary>
        public ReadOnlyCollection<ZpkArchiveEntry> Entries
        {
            get { return m_entries; }
        }

        /// <summary>
        /// Retrieves a wrapper for the specified entry in the zpk archive.
        /// </summary>
        /// <param name="entryName">A path, relative to the root of the 
        /// archive, that identifies the entry to retrieve.</param>
        /// <returns>A wrapper for the specified entry in the archive; null
        /// if the entry does not exist in the archive.</returns>
        public ZpkArchiveEntry GetEntry(string entryName)
        {
            foreach (ZpkArchiveEntry entry in m_entries)
            {
                if (entry.FullName == entryName)
                    return entry;
            }
            return null;
        }

        public string FileName
        {
            get { return m_filename; }
        }

        internal IEnumerable<Stream> GetEntryData(ZpkArchiveEntry entry)
        {
            // Opens the file. We do this each time so that we can read
            // two entries simulnaneously.
            Stream fileStream = File.OpenRead(m_filename);

            // Decrypt and then decompress the entry data chunk by chunk. For
            // each chunk, we yield return a stream from which the chunk's
            // data can be read.
            long chunkOffset = entry.CompressedOffset;
            long remainingLength = entry.CompressedLength;
            while (remainingLength > 0)
            {
                Stream decompressedStream = null;
                try
                {
                    // Seek to the chunk data region.
                    fileStream.Seek(chunkOffset, SeekOrigin.Begin);

                    // Create a decrypted stream. This is the whole entry
                    // stream if the entry is not encrypted; otherwise, it
                    // is a chunk with a pair of length prefixes.
                    Stream encryptedStream, decryptedStream;
                    if (entry.EncryptionMethod == ZpkEncryptionMethod.None)
                    {
                        // Create a limited length stream to wrap the given region.
                        encryptedStream = new LimitedLengthStream(
                            fileStream, entry.CompressedLength, true);
                        decryptedStream = encryptedStream;
                        chunkOffset += encryptedStream.Length;
                        remainingLength -= decryptedStream.Length;
                    }
                    else
                    {
                        int decryptedLength = fileStream.ReadInt32(Endian.Big);
                        int encryptedLength = fileStream.ReadInt32(Endian.Little);
                        encryptedStream = new LimitedLengthStream(
                            fileStream, encryptedLength, true);
                        chunkOffset += (8 + encryptedStream.Length);
                        // Update the number of bytes that remain to be read. 
                        // This number can be negative if entry.CompressedLength
                        // does not match the total length of the chunks.
                        remainingLength -= (8 + decryptedLength);

                        ICryptoTransform cipher;
                        switch (entry.EncryptionMethod)
                        {
                            case ZpkEncryptionMethod.TEA:
                                cipher = tea;
                                break;
                            case ZpkEncryptionMethod.XTEA:
                                cipher = xtea;
                                break;
                            default:
                                throw new NotSupportedException("EncryptionMethod");
                        }
                        decryptedStream = new LimitedLengthStream(
                            new CryptoStream(encryptedStream, cipher, CryptoStreamMode.Read),
                            decryptedLength,
                            false);
                    }

                    // Create a decompressed stream.
                    switch (entry.CompressionMethod)
                    {
                        case ZpkCompressionMethod.None:
                            decompressedStream = decryptedStream;
                            break;
                        case ZpkCompressionMethod.LZMA86:
                            // Check byte.
                            byte[] b = new byte[100];
                            decryptedStream.ReadFull(b);
                            if (decryptedStream.ReadByte() != 0)
                                throw new NotSupportedException("BCJ filter not implemented.");
                            decompressedStream = new LzmaStream(
                                decryptedStream, CompressionMode.Decompress);
                            break;
                        case ZpkCompressionMethod.LZSS:
                        default:
                            throw new NotSupportedException();
                    }
                }
                catch (Exception)
                {
                    fileStream.Dispose();
                    throw;
                }
                yield return decompressedStream;
            }

            // Close the file stream.
            fileStream.Dispose();

#if false
                
                // Uncompress the block.
                switch (st.compression)
                {
                    case ZpkCompression.None:

                        if (iOutput + streamlen > outlen)
                            throw new ZpkFormatException();

                        Buffer.BlockCopy(inbuf, (int)i, outbuf, (int)iOutput, (int)streamlen);
                        iOutput += (int)streamlen;
                        break;

                    case ZpkCompression.Lzss:

                        if (streamlen < 4)
                            throw new ZpkFormatException();
                        uint lzssLen = BigEndian.ToUInt32(inbuf, i);
                        if (lzssLen + 4 > streamlen)
                            throw new ZpkFormatException();
                        try
                        {
                            iOutput += LzssDecoder.Decode(
                                inbuf, i + 4, (int)lzssLen, outbuf, iOutput, (int)(outlen - iOutput));
                        }
                        catch (Exception)
                        {
                            throw new ZpkFormatException();
                        }
                        break;

                    case ZpkCompression.Lzma:

                        if (streamlen < 4)
                            throw new ZpkFormatException();
                        uint lzmaLen = BigEndian.ToUInt32(inbuf, i);
                        if (lzmaLen + 4 > streamlen)
                            throw new ZpkFormatException();

                        try
                        {
                            iOutput += LzmaDecoder.Decode(
                                inbuf, i + 4, (int)lzmaLen, outbuf, iOutput, (int)(outlen - iOutput));
                        }
                        catch (LzmaException)
                        {
                            throw new ZpkFormatException();
                        }
                        break;

                    default:
                        throw new ZpkNotSupportedException(
                            "不支持压缩格式 " + ((int)st.compression).ToString());
                }

                // Advance the input pointer.
                i += (int)blocklen;
            }

            // Make sure we write exactly the number of bytes we expect.
            if (iOutput != outlen)
                throw new ZpkFormatException();

            return outbuf;
        }
#endif
        }
#if false

        public int Version
        {
            get { return m_version; }
        }

        public byte[] ExtractEntry(int index)
        {
            ZpkEntryStat st = m_entries[index];
            byte[] inbuf = new byte[st.packed_size];
            byte[] outbuf = new byte[st.size];

            // Read the compressed data into memory.
            using (FileStream fs = new FileStream(m_filename, FileMode.Open, FileAccess.Read, FileShare.Read))
            {
                fs.Seek(st.packed_offset, SeekOrigin.Begin);
                ReadFull(fs, inbuf, 0, (int)st.packed_size);
            }

            // Extract the compressed file entry block by block.
            uint inlen = st.packed_size, outlen = st.size;
            int i = 0, iOutput = 0;
            while (iOutput < outlen)
            {

                uint streamlen = st.packed_size;    // Length of the real input stream
                uint blocklen = streamlen;          // Length of the encapsulating stream

                // Decrypt the block, such that, after decryption:
                //   i := index to start of stream
                //   streamlen := real length of input stream
                //   blocklen := number of bytes to advance i to next block
                if (st.encryption != ZpkEncryption.None)
                {

                    // Read streamlen and blocklen.
                    if (i + 8 > inlen)
                        throw new ZpkFormatException();
                    streamlen = BigEndian.ToUInt32(inbuf, i);
                    blocklen = LittleEndian.ToUInt32(inbuf, i + 4);
                    i += 8;
                    if (i + blocklen > inlen || streamlen > blocklen)
                        throw new ZpkFormatException();

                    // Decrypt the blocks.
                    switch (st.encryption)
                    {
                        case ZpkEncryption.TEA:
                            m_tea.Decrypt(inbuf, i, blocklen);
                            break;
                        case ZpkEncryption.XTEA:
                            m_xtea.Decrypt(inbuf, i, blocklen);
                            break;
                        default:
                            throw new ZpkNotSupportedException(
                                "不支持加密格式 " + ((int)st.encryption).ToString());
                    }
                }

                // Uncompress the block.
                switch (st.compression)
                {
                    case ZpkCompression.None:

                        if (iOutput + streamlen > outlen)
                            throw new ZpkFormatException();

                        Buffer.BlockCopy(inbuf, (int)i, outbuf, (int)iOutput, (int)streamlen);
                        iOutput += (int)streamlen;
                        break;

                    case ZpkCompression.Lzss:

                        if (streamlen < 4)
                            throw new ZpkFormatException();
                        uint lzssLen = BigEndian.ToUInt32(inbuf, i);
                        if (lzssLen + 4 > streamlen)
                            throw new ZpkFormatException();
                        try
                        {
                            iOutput += LzssDecoder.Decode(
                                inbuf, i + 4, (int)lzssLen, outbuf, iOutput, (int)(outlen - iOutput));
                        }
                        catch (Exception)
                        {
                            throw new ZpkFormatException();
                        }
                        break;

                    case ZpkCompression.Lzma:

                        if (streamlen < 4)
                            throw new ZpkFormatException();
                        uint lzmaLen = BigEndian.ToUInt32(inbuf, i);
                        if (lzmaLen + 4 > streamlen)
                            throw new ZpkFormatException();

                        try
                        {
                            iOutput += LzmaDecoder.Decode(
                                inbuf, i + 4, (int)lzmaLen, outbuf, iOutput, (int)(outlen - iOutput));
                        }
                        catch (LzmaException)
                        {
                            throw new ZpkFormatException();
                        }
                        break;

                    default:
                        throw new ZpkNotSupportedException(
                            "不支持压缩格式 " + ((int)st.compression).ToString());
                }

                // Advance the input pointer.
                i += (int)blocklen;
            }

            // Make sure we write exactly the number of bytes we expect.
            if (iOutput != outlen)
                throw new ZpkFormatException();

            return outbuf;
        }
#endif

#if false
    public class ZpkArchive
    {
        private List<ZpkVolume> m_volumeList;
        private List<ZpkEntry> m_entryList;
        private Dictionary<string, int> m_entryMap;

        public ZpkArchive(string filename, bool loadPatches)
        {
            // Load the main volume.
            ZpkVolume volMain = new ZpkVolume(filename, 0);

            // Allocate members.
            m_volumeList = new List<ZpkVolume>();
            m_entryList = new List<ZpkEntry>(volMain.EntryCount);
            m_entryMap = new Dictionary<string, int>(volMain.EntryCount);

            // Add the main volume.
            AddVolume(volMain);

            // Search for patches of the archive in the same directory.
            if (loadPatches) {
                DirectoryInfo dirInfo = new DirectoryInfo(Path.GetDirectoryName(filename));
                string name = Path.GetFileName(filename);
                foreach (FileInfo fi in dirInfo.GetFiles()) {
                    if (!fi.Name.StartsWith(name + ".", StringComparison.InvariantCultureIgnoreCase))
                        continue;
                    string ext = fi.Name.Substring(name.Length + 1);
                    int version;
                    if (Int32.TryParse(ext, out version)) {
                        try {
                            ZpkVolume vol = new ZpkVolume(fi.FullName, version);
                            AddVolume(vol);
                        } catch (ZpkFormatException) {
                            continue;
                        }
                    }
                }
            }
        }

        public void AddVolume(ZpkVolume vol)
        {
            m_volumeList.Add(vol);

            // Create entries for items in this volume.
            for (int i = 0; i < vol.EntryCount; i++) {
                ZpkEntry ent = new ZpkEntry(vol, i);
                string key = ent.Name.ToLowerInvariant();
                int index;
                if (m_entryMap.TryGetValue(key, out index)) {
                    if (m_entryList[index].Version < ent.Version) 
                        m_entryList[index] = ent;
                } else {
                    index = m_entryList.Count;
                    m_entryList.Add(ent);
                    m_entryMap.Add(key, index);
                }
            }
        }

        public ZpkVolume[] Volumes
        {
            get { return m_volumeList.ToArray(); }
        }

        public int EntryCount
        {
            get { return m_entryList.Count; }
        }

        public ZpkEntry GetEntry(string name)
        {
            int index;
            if (m_entryMap.TryGetValue(name.ToLowerInvariant(), out index)) {
                return m_entryList[index];
            } else {
                return null;
            }
        }

        public ZpkEntry GetEntry(int index)
        {
            return m_entryList[index];
        }
    
    }
#endif
    }

    /// <summary>
    /// Contains information about a file entry in a ZPK volume.
    /// </summary>
    public class ZpkArchiveEntry
    {
        /// <summary>
        /// Gets the ZpkArchive that contains this entry.
        /// </summary>
        public ZpkArchive Archive { get; internal set; }

        /// <summary>
        /// Gets the relative path of the entry in the zpk archive.
        /// </summary>
        public string FullName { get; internal set; }

        /// <summary>
        /// Gets the uncompressed size of the entry in the zpk archive.
        /// </summary>
        public uint Length { get; internal set; }

        internal uint CompressedOffset { get; set; }

        /// <summary>
        /// Gets the compressed size of the entry in the zpk archive.
        /// </summary>
        public uint CompressedLength { get; internal set; }

        /// <summary>
        /// Gets the encryption method used to store the entry in the zpk
        /// archive.
        /// </summary>
        public ZpkEncryptionMethod EncryptionMethod { get; internal set; }

        /// <summary>
        /// Gets the compression method used to store the entry in the zpk
        /// archive.
        /// </summary>
        public ZpkCompressionMethod CompressionMethod { get; internal set; }

        public override string ToString()
        {
            return this.FullName;
        }

        /// <summary>
        /// Opens the entry from the zpk archive for reading.
        /// </summary>
        /// <returns>
        /// The stream that represents the contents of the entry.
        /// </returns>
        public Stream Open()
        {
            return new ConcatenatedStream(this.Archive.GetEntryData(this));
        }

#if false
        public int Version
        {
            get { return m_volume.Version; }
        }

        public byte[] Extract()
        {
            return m_volume.ExtractEntry(m_index);
        }
#endif
    }

    /// <summary>
    /// Defines the encryption method used to store a file entry in a ZPK
    /// archive.
    /// </summary>
    public enum ZpkEncryptionMethod
    {
        /// <summary>
        /// The entry is not encrypted.
        /// </summary>
        None = 0,

        /// <summary>
        /// The entry is encrypted using Tiny Encryption Algorithm.
        /// </summary>
        TEA = 1,

        /// <summary>
        /// The entry is encrypted using Extended Tiny Encryption Algorithm.
        /// </summary>
        XTEA = 2,
    }

    /// <summary>
    /// Defines the compression method used to store a file entry in a ZPK
    /// archive.
    /// </summary>
    public enum ZpkCompressionMethod
    {
        /// <summary>
        /// The entry is not compressed.
        /// </summary>
        None = 0,

        /// <summary>
        /// The entry is compressed using LZSS algorithm.
        /// </summary>
        LZSS = 1,

        /// <summary>
        /// The entry is compressed using LZMA86 algorithm.
        /// </summary>
        LZMA86 = 2,
    }
}
